import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
import pandas as pd
from dataset.vocab import Vocab, tokenizer
import re

df = pd.read_csv("./data/translate_chinese2english.csv")
data = df.to_numpy()

# 加入<bos> 与<eos>
en_data = []
for item in data[:, 1]:
    en_data.append("<bos> " + item + " <eos>")

en_tokens = tokenizer(en_data, mode="word")
zh_tokens = tokenizer(data[:, 0], mode="char")
en_vocab = Vocab(en_tokens, 0, retired_tokens=['<pad>'])  # <pad> == 0 此时<pad> == <unk>
zh_vocab = Vocab(zh_tokens, 0, retired_tokens=['<pad>'])
en_idx = [torch.tensor(en_vocab.to_idx(line)) for line in en_tokens]
zh_idx = [torch.tensor(zh_vocab.to_idx(line)) for line in zh_tokens]
en_valid_len = [len(line) for line in en_tokens]
zh_valid_len = [len(line) for line in zh_tokens]


def get_valid_len():
    """
    :return: 每个句子的有效长度
    """
    return zh_valid_len, en_valid_len


def generate_vocab():
    return zh_vocab, en_vocab


class TranslateDataset(Dataset):
    def __init__(self, sentence, translate_sentence, zh_valid_len, en_valid_len):
        super().__init__()
        self.sentence = sentence
        self.translate_sentence = translate_sentence
        self.zh_valid_len = zh_valid_len
        self.en_valid_len = en_valid_len

    def __len__(self):
        return len(self.sentence)

    def __getitem__(self, index):
        return self.sentence[index], self.translate_sentence[index], self.zh_valid_len[index], self.en_valid_len[index]


# 将一个此次中数据全部对齐
def collate_fn(batch):
    zh_inputs, en_inputs, tmp1, tmp2 = zip(*batch)
    zh_pad = pad_sequence(zh_inputs, batch_first=True, padding_value=zh_vocab.to_idx("<pad>"))
    en_pad = pad_sequence(en_inputs, batch_first=True, padding_value=en_vocab.to_idx("<pad>"))
    return zh_pad, en_pad, tmp1, tmp2


def generate_loader(batch_size=20):
    dataset = TranslateDataset(zh_idx, en_idx, zh_valid_len, en_valid_len)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
    return dataloader
